import os
from fastapi import FastAPI, Body
from pydantic import BaseModel
import gradio as gr
import json
import gc
import torch
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
from modules import scripts, shared, script_callbacks
from modules.paths import models_path
from modules.devices import device, torch_gc, cpu
from modules.safe import unsafe_torch_load, load
from collections import OrderedDict
from modules.api.api import encode_pil_to_base64, decode_base64_to_image
from PIL import Image
import numpy as np
from typing import Any, Optional, List


scripts_sam_model_dir = os.path.join(scripts.basedir(), "models")
sd_sam_model_dir = os.path.join(models_path, "sam")
sam_model_dir = sd_sam_model_dir if os.path.exists(sd_sam_model_dir) else scripts_sam_model_dir
sam_model_list = [f for f in os.listdir(sam_model_dir) if os.path.isfile(os.path.join(sam_model_dir, f)) and f.split('.')[-1] != 'txt']
sam_device = device
sam_model_cache = OrderedDict()

def decode_to_pil(image):
    if os.path.exists(image):
        return Image.open(image)
    elif type(image) is str:
        return decode_base64_to_image(image)
    elif type(image) is Image.Image:
        return image
    elif type(image) is np.ndarray:
        return Image.fromarray(image)
    else:
        Exception("Not an image")


def encode_to_base64(image):
    if type(image) is str:
        return image
    elif type(image) is Image.Image:
        return encode_pil_to_base64(image).decode()
    elif type(image) is np.ndarray:
        pil = Image.fromarray(image)
        return encode_pil_to_base64(pil).decode()
    else:
        Exception("Invalid type")

def sd_webui_sam_api(_: gr.Blocks, app: FastAPI):
    @app.get("/samapi/heartbeat")
    async def heartbeat():
        return {
            "msg": "Success!"
        }

    @app.get("/samapi/sam-model", description='查询可以用的SAM model')
    async def api_sam_model() -> List[str]:
        return sam_model_list



    def load_sam_model(sam_checkpoint):
        model_type = sam_checkpoint.split('.')[0]
        if 'hq' not in model_type and 'mobile' not in model_type:
            model_type = '_'.join(model_type.split('_')[1:-1])
        sam_checkpoint_path = os.path.join(sam_model_dir, sam_checkpoint)
        torch.load = unsafe_torch_load
        sam = sam_model_registry[model_type](checkpoint=sam_checkpoint_path)
        sam.to(device=sam_device)
        sam.eval()
        torch.load = load
        return sam

    def clear_sam_cache():
        sam_model_cache.clear()
        gc.collect()
        torch_gc()

    def clear_cache():
        clear_sam_cache()

    def garbage_collect(sam):
        # if shared.cmd_opts.lowvram:
        #     sam.to(cpu)
        gc.collect()
        torch_gc()

    def init_sam_model(sam_model_name):
        print(f"Initializing SAM model {sam_model_name}")
        if sam_model_name in sam_model_cache:
            sam = sam_model_cache[sam_model_name]
            # if shared.cmd_opts.lowvram or (str(sam_device) not in str(sam.device)):
            #     sam.to(device=sam_device)
            return sam
        elif sam_model_name in sam_model_list:
            clear_sam_cache()
            sam_model_cache[sam_model_name] = load_sam_model(sam_model_name)
            return sam_model_cache[sam_model_name]
        else:
            raise Exception(
                f"{sam_model_name} not found, please download model to models.")



    # sam_model_name sam_vit_h_4b8939.pth,sam_vit_l_0b3195.pth,sam_vit_b_01ec64.pth
    # image -- base64编码的图片
    # output_mode :"binary_mask",
    #             "uncompressed_rle",
    #             "coco_rle"
    def sam_generate_all_mask(sam_model_name,output_mode,image):
        sam = init_sam_model(sam_model_name)
        mask_generator = SamAutomaticMaskGenerator(sam, output_mode=output_mode)
        image = decode_to_pil(image)
        # input_image_np = np.array(image)
        img_np = np.array(image.convert("RGB"))
        masks = mask_generator.generate(img_np)
        return json.dumps(masks, separators=(",", ":"))

    class SamPredictRequest(BaseModel):
        sam_model_name: str
        output_mode: str
        input_image: str

    class SamPredictRequestWithImg(BaseModel):
        input_image: str


    @app.post("/samapi/auto-sam-vit-b", description='sam图片使用vit-b模型')
    async def api_sam_model(payload: SamPredictRequestWithImg = Body(...)) -> Any:
        sam_model_name="sam_vit_b_01ec64.pth"
        output_mode="uncompressed_rle"
        return sam_generate_all_mask(sam_model_name,output_mode,payload.input_image)

    @app.post("/samapi/auto-sam-vit-l", description='sam图片使用vit-l模型')
    async def api_sam_model(payload: SamPredictRequestWithImg = Body(...)) -> Any:
        sam_model_name = "sam_vit_l_0b3195.pth"
        output_mode = "uncompressed_rle"
        return sam_generate_all_mask(sam_model_name, output_mode, payload.input_image)

    @app.post("/samapi/auto-sam-vit-h", description='sam图片使用vit-h模型')
    async def api_sam_model(payload: SamPredictRequestWithImg = Body(...)) -> Any:
        sam_model_name = "sam_vit_h_4b8939.pth"
        output_mode = "uncompressed_rle"
        return sam_generate_all_mask(sam_model_name, output_mode, payload.input_image)

    @app.post("/samapi/auto-sam", description='自动sam图片')
    async def api_sam_model(payload: SamPredictRequest = Body(...)) -> Any:
        return sam_generate_all_mask(payload.sam_model_name,payload.output_mode,payload.input_image)


try:
    import modules.script_callbacks as script_callbacks
    script_callbacks.on_app_started(sd_webui_sam_api)
except:
    print("SAM API failed to initialize")